package com.castlzl.random.core.tinyid.zookeeper;

import com.castlzl.random.core.tinyid.TinyIdGenerator.Segment;
import com.castlzl.random.core.tinyid.TinyIdSegmentService;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.locks.InterProcessMutex;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.data.Stat;
import org.springframework.stereotype.Service;

import java.nio.charset.StandardCharsets;
import java.util.concurrent.TimeUnit;

/**
 * ZooKeeper实现的号段分配服务
 * 使用ZooKeeper的分布式锁保证号段分配的安全性
 * @author leiziliang
 */
@Slf4j
@Service
public class ZookeeperTinyIdSegmentServiceImpl implements TinyIdSegmentService {

    private final CuratorFramework client;
    private final long step;
    private final String rootPath;
    private static final int LOCK_TIMEOUT_MS = 5000;

    private static final String PREFIX = "/";

    public ZookeeperTinyIdSegmentServiceImpl(CuratorFramework client, long step, String rootPath) {
        this.client = client;
        this.step = step;
        if (!rootPath.startsWith(PREFIX)) {
            this.rootPath = PREFIX+rootPath;
        }else{
            this.rootPath = rootPath;
        }
        initZookeeperPath();
    }

    private void initZookeeperPath() {
        try {
            Stat stat = client.checkExists().forPath(rootPath);
            if (stat == null) {
                log.debug("Creating zookeeper path: {}", rootPath);
                createPersistentPath(rootPath);
                log.info("Successfully created zookeeper path: {}", rootPath);
            } else {
                log.debug("Zookeeper path already exists: {},{}", rootPath, stat.getDataLength());
            }
            
            // 确保业务标签目录存在
            String bizTagsPath = rootPath + "/biztags";
            if (client.checkExists().forPath(bizTagsPath) == null) {
                createPersistentPath(bizTagsPath);
                log.info("Successfully created biztags path: {}", bizTagsPath);
            }
        } catch (Exception e) {
            log.error("Failed to init zookeeper path: {}", rootPath, e);
            throw new RuntimeException("Failed to init zookeeper path: " + rootPath, e);
        }
    }

    private void createPersistentPath(String path) throws Exception {
        client.create()
                .creatingParentsIfNeeded()
                .withMode(CreateMode.PERSISTENT)
                .forPath(path);
    }

    @Override
    public Segment getNextSegment(String bizTag) {
        String bizPath = rootPath + "/biztags/" + bizTag;
        String lockPath = rootPath + "/lock/" + bizTag;
        InterProcessMutex lock = new InterProcessMutex(client, lockPath);
        boolean locked = false;
    
        try {
            // 获取分布式锁
            if (!lock.acquire(LOCK_TIMEOUT_MS, TimeUnit.MILLISECONDS)) {
                throw new RuntimeException("Failed to acquire lock for: " + bizTag);
            }
            locked = true;
    
            long currentValue;
            // 原子化处理节点创建和数据获取
            try {
                byte[] data = client.getData().forPath(bizPath);
                currentValue = parseSegmentValue(data, bizPath);
            } catch (KeeperException.NoNodeException ex) {
                try {
                    // 创建持久化节点，初始值为0
                    client.create()
                         .creatingParentsIfNeeded()
                         .withMode(CreateMode.PERSISTENT)
                         .forPath(bizPath, "0".getBytes(StandardCharsets.UTF_8));
                    currentValue = 0;
                    log.info("Created new bizTag node: {} with initial value: 0", bizPath);
                } catch (KeeperException.NodeExistsException e) {
                    // 处理并发创建的情况
                    byte[] data = client.getData().forPath(bizPath);
                    currentValue = parseSegmentValue(data, bizPath);
                    log.debug("Using existing bizTag node: {} with value: {}", bizPath, currentValue);
                }
            }
    
            // 更新并持久化新值
            long newValue = currentValue + step;
            client.setData().forPath(bizPath, String.valueOf(newValue).getBytes(StandardCharsets.UTF_8));
            log.debug("Updated bizTag node: {} with new value: {}", bizPath, newValue);
    
            return new Segment(currentValue + 1, newValue);
        } catch (Exception e) {
            log.error("Failed to get segment for bizTag: {}", bizTag, e);
            throw new RuntimeException("Failed to get segment for: " + bizTag, e);
        } finally {
            if (locked) {
                try {
                    lock.release();
                } catch (Exception e) {
                    log.error("Failed to release lock for bizTag: {}", bizTag, e);
                }
            }
        }
    }
    
    private long parseSegmentValue(byte[] data, String path) {
        try {
            return Long.parseLong(new String(data, StandardCharsets.UTF_8));
        } catch (NumberFormatException e) {
            throw new RuntimeException("Invalid number format in node: " + path, e);
        }
    }
}